# 了解各个函数的功能。
from androguard.misc import AnalyzeAPK
from collections import defaultdict
# coding=utf-8
import csv
import os
import networkx as nx
import re
from feature.LibRadar.literadar import LibRadarLite
# from feature.Andro.permission import Permission


"""----------------------------------------------UTILS METHOD-------------------------------------------------------------"""
def write_csv(opcode, f, id):  # 写csv
    csv_write = csv.writer(f)
    data_row = [id]
    for op in opcode.keys():
        data_row.append(opcode[op])
    csv_write.writerow(data_row)

def create_csv(smali_opcode, path):  # 创建csv
    f = open(path, 'w+', newline='')
    csv_write = csv.writer(f)
    csv_head = ['id'] + smali_opcode
    csv_write.writerow(csv_head)
    return f

def get_from_csv_gml(filename):  # 从API.csv获取每个API对应的权限mapping
    per_value = {}
    with open(filename, "r") as csvFile:
        reader = csv.reader(csvFile)
        for item in reader:
            if reader.line_num == 1:
                continue
            name = item[0]
            per = item[1]
            # value[name] = functionname
            if name not in per_value.keys():
                per_value[name] = [per]
            else:
                per_value[name].append(per)
    return per_value

def get_external(nodeid, G):
    return G.nodes[nodeid]['external']

def n_neighbor(node, graph, hop=1):  # 节点的一跳邻居
    ego_graph = nx.ego_graph(graph, node, radius=hop, undirected=True)
    nodes = ego_graph.nodes
    return nodes

def get_label(node_id, G):
    return G.nodes[node_id]['label']
"""--------------------------------------------------------------------------------------------------------------------"""

# cppermission
def get_permission(path):
    filename = path
    permission = {}
    with open(filename) as f:
        content = f.readline().strip('\n')
        while content:
            cons = content.split(' ')
            if cons[0] not in permission:
                permission[cons[0]] = set()
            permission[cons[0]].add((cons[1], 'Permission:' + cons[2]))
            content = f.readline().strip('\n')
    return permission

def df_from_G(G):
    import pandas as pd
    df = pd.DataFrame(G.nodes(data=True))
    try:
        attr = df[1].apply(pd.Series)
    except KeyError:
        return False
    node_attr = pd.concat([df[0], attr], axis=1)
    node_attr = node_attr.rename(columns={0: 'id'})
    return node_attr

def node2function(s):  # 根据节点名称获取函数名称
    first = s.find(' ')
    s = s[first + 1:]
    if s.find('@ ') < 0:
        first = s.rfind('>')
        s = s[:first]
    else:
        first = s.find('[access_flags')
        s = s[:first - 1]
    return s

def getMethod2NodeMap(G):
    method2nodeMap = {}
    try:
        node_attr = df_from_G(G)
        labels = node_attr.label
        ids = node_attr.id
    except Exception:
        return method2nodeMap
    i = 0
    pattern = re.compile(r'&#(.+?);')
    while i < len(ids):
        nodeid = ids.get(i)
        label = labels.get(i)
        function = node2function(label)
        rt = pattern.findall(function)
        for r in rt:
            function.replace("&#%s;"%r, chr(int(r)))
        method = function.replace(";->", "; ").replace("(", " (")
        method2nodeMap.update({method: (nodeid, function)})
        i = i + 1
    return method2nodeMap

def get_smaliOpcode(filename):  # read all smali opcode list
        opcode = list()
        with open(filename, 'r') as fileObject:
            lines = fileObject.readlines()
        for line in lines:
            opcode.append(line.rstrip('\n'))
        return opcode

"""----------------------------------------------PERMISSOIN METHOD-------------------------------------------------------------"""
def is_in_funcList(funcList, t):  # 节点是否再函数列表中
    for f in funcList:
        # 假如在，就是true
        if t.find(f) >= 0:
            return True
    return False

def get_nodeid_label(G, function):  # 通过方法（类+函数） 获取id和label
    if type(function) == int:
        return function, G.nodes[function]['label']
    nodes = G.nodes
    for node in nodes:
        label = G.nodes[node]['label']
        if label.find(function) >= 0:
            return node, label
    return "", ""

def getclass(functionname):
    index = functionname.find(';->')
    return functionname[len('<analysis.MethodAnalysis L'):index]

# 将方法(外部) 映射为  原来的名称: 父类名称 ,  id
def get_mappings(function, classname):
    # 传入方法和类名。
    mappings = {}
    nodeid, nodelabel = get_nodeid_label(G, function)
    #  id + label
    if nodeid == "":
        return mappings
    classname = 'L' + classname.replace('<analysis.MethodAnalysis ', '') + ';'
    funcList = []
    try:
        # 获得某个类所有的方法
        tmp = class_functions[classname]

        for t in tmp:
            # 获得每个方法的全程 funcList
            funcList.append(method2nodeMap[t][1])
    except KeyError:
        funcList = []
    
    # funcList 某个c中，全部的方法名称 ['Lcom/qihoo/util/Configuration; <clinit> ()V', 'Lcom/qihoo/util/Configuration; <init> ()V']
    t_id = nodeid
    t = nodelabel
    # 判断传入的类是否为外部类。
    external = get_external(nodeid, G)
    # 假如是external(不重要吧。)
    if external == 1:
        t_class = classname
        # 假如不在则执行
        if not is_in_funcList(funcList, t):  # t is in super class or system function
            if t_class in super_dic.keys():
                super_class = super_dic[t_class]
                while True:
                    new_label = t.replace(t_class, super_class)
                    try:
                        super_funcList = class_functions[super_class]
                        if is_in_funcList(super_funcList, new_label):
                            mappings[t] = (new_label, t_id)
                            break
                        else:
                            t_class = super_class
                            super_class = super_dic[t_class]
                    except KeyError:
                        mappings[t] = (new_label, t_id)
                        break
    return mappings
    
def substitude():  # 替换成子类
    functions = node_attr.label
    ids = node_attr.id
    # start method connect run method
    for c in class_functions:
        if c in implement_dic:
            super_c = implement_dic[c]
        else:
            super_c = ""
        # 如果在这个replacemap中。
        if super_c in replacemap:
            index = 0 
            while index < len(ids):
                func = functions.get(index)
                if func.find(c + "-><init>(L") >= 0:
                    left = func.find(";-><init>(L") + len(";-><init>(L")
                    right = func.find(";", left)
                    baseclass = func[left: right]
                    # c与baseclass有什么区别？？？？？？？？
                    index2 = 0
                    func_list = replacemap[super_c]
                    while index2 < len(ids):
                        func_tmp = functions.get(index2)
                        for ftmp in func_list:
                            if func_tmp.find(c + "->" + ftmp) >= 0:
                                class2runinit[baseclass].update({ids.get(index2): super_c[1:-1]})
                        index2 = index2 + 1
                    break
                index = index + 1
    # print(class2runinit)   {'android/content/Context': {4: 'java/lang/Runnable'}})
    mappings = {}
    index = 0
    # 将外部类（名字可能是内部类的方法）映射成  内部类+函数： 外部内+函数  ， 编号ID
    #  Lcom/qihoo/util/StubApp1520421995;->getApplicationContext()Landroid/content/Context;>': ('<analysis.MethodAnalysis Landroid/app/Application;->getApplicationContext()Landroid/content/Context;>', 59)
    while index < len(ids):
        # 获取 当前index 的 类；->方法 。
        label = functions.get(index)
        # 获取 当前index 的 类名 com/qihoo/util/StubApplication
        classname = getclass(label)
        mappings.update(get_mappings(label, classname))
        index = index + 1
        
    # 将外部类（名字可能是内部类的方法）映射成  内部类+函数： 外部内+函数  ， 编号ID
    class_functions_copy = class_functions.copy()
    for classname in class_functions_copy:
        for function in class_functions_copy[classname]:
            label = method2nodeMap[function][1]
            mappings.update(get_mappings(label, classname))
    # print(mappings)
    # print(super_dic)
    res = {}
    for function in mappings:
        super_function = mappings[function][0]
        for func in per_map:
            if super_function.find(func) >= 0:
                res[mappings[function][1]] = []
                for p in per_map[func]:
                    res[mappings[function][1]].append(p)
    return res  # 所有被替换成子类的敏感API


def deal_node(nodeid):
    permission_node = {}
    # 获取其后继节点。
    targets = G.successors(nodeid)
    for t in targets:
        # 获取邻居的lable
        label = get_label(t,G)
        for k in method_dic.keys():
            # k=lable
            # 假如在邻居中发现method_dic.keys
            if label.find(k) >= 0:
                for ki in method_dic[k]:
                    if ki not in permission_node:
                        permission_node[ki] = set()
                        # permission_node[W] = 1,2,3,4
                        # permission_node[R] = 2,1230123,12315...
                    permission_node[ki].add(nodeid)
    return permission_node

# Permission包含api和contentprovider的部分
# Cursor query(Uri uri, String[] projection, String selection, String[] selectionArgs, String sortOrder)
query_method = "Landroid/content/ContentResolver;->query(Landroid/net/Uri; [Ljava/lang/String; Ljava/lang/String; [Ljava/lang/String; Ljava/lang/String;)Landroid/database/Cursor;"
# Uri insert(Uri uri, ContentValues values)
insert_method = "Landroid/content/ContentResolver;->insert(Landroid/net/Uri; Landroid/content/ContentValues;)Landroid/net/Uri;"
# int delete(Uri uri, String selection, String[] selectionArgs)
delete_method = "Landroid/content/ContentResolver;->delete(Landroid/net/Uri; Ljava/lang/String; [Ljava/lang/String;)I"
# int update(Uri uri, ContentValues values, String selection, String[] selectionArgs)
update_method = "Landroid/content/ContentResolver;->update(Landroid/net/Uri; Landroid/content/ContentValues; Ljava/lang/String; [Ljava/lang/String;)I"

method_dic = {query_method: ['R'], insert_method: ['W'], delete_method: ['W'], update_method: ['W', 'R']}

"""--------------------------------------------------------------------------------------------------------------------"""



"""----------------------------------------------FILE PATH-------------------------------------------------------------"""

smali_opcode=get_smaliOpcode('/home/chenjiaxiong/jx_code/MsDroid-JXC/feature/Andro/Data/smaliOpcode.txt')
cppermission = get_permission('/home/chenjiaxiong/jx_code/MsDroid-JXC/feature/Andro/Data/APIs/all_cp.txt')
permission = []
with open('/home/chenjiaxiong/jx_code/MsDroid-JXC/feature/Andro/Data/head.txt') as f:
    permission = eval(f.read())
opcodeFilename ='/home/chenjiaxiong/jx_code/android_sample/csv/opcode.csv'
permissionFilename ='/home/chenjiaxiong/jx_code/android_sample/csv/permisson.csv'
tplFilename = '/home/chenjiaxiong/jx_code/android_sample/csv/tpl.csv'
CGfilename = '/home/chenjiaxiong/jx_code/android_sample/graph/98520bc4ea00c4566e46d533e2f263d0.gml'
opcodeFile = create_csv(smali_opcode,opcodeFilename)
replacemap = {'Landroid/os/AsyncTask;': ['onPreExecute', 'doInBackground'],
              'Landroid/os/Handler;': ['handleMessage'], 'Ljava/lang/Runnable;': ['run']}
"""--------------------------------------------------------------------------------------------------------------------"""


"""----------------------------------------------APK GRAPH-------------------------------------------------------------"""
# 355f506a84bf3034dc8b4474a7f1c8ef巨慢巨慢巨慢巨慢！！！
# 5dd224d15f6af17dae412773cbc0ff16处理很慢
# b323e5aa2926181300a9faf253224d47很慢很慢，要睡着了。
# 9d4c00525f089ce0f589fadcffc7d01b 有点慢。

# 3d3ac09d1f058e5d8f1c54ecbc9a6b5a很快，有点参考价值。
# 6e91ff990b50cc946812c17a2035c983很快，但是没有参考价值。
# 9854adb54adcc325b82bd3e3459327eb 很快。 双空。

apkfile = '/home/chenjiaxiong/jx_code/android_sample/apk/98520bc4ea00c4566e46d533e2f263d0.apk'  # APK 文件路径
a, d, dx = AnalyzeAPK(apkfile)
G = dx.get_call_graph()

nx.write_gml(G, CGfilename, stringizer=str)  # save the call graph

print('已经生成图啦.')
G = nx.read_gml(CGfilename, label='id')
"""--------------------------------------------------------------------------------------------------------------------"""

"""----------------------------------------------APK ANALYSIS----------------------------------------------------------"""
# 每个class 有 几种方法.
class_functions = defaultdict(list)  # mappings of class and its functions
# class和他的父类class
super_dic = {}  # mappings of class and its superclass(for class replacement)
implement_dic = {}
method2nodeMap = getMethod2NodeMap(G)

# 生成每个method + csv
for classes in dx.get_classes():  # all class
    class_name = str(classes.get_class().get_name())
    """
    判断它的父类,假如不是 "Ljava/lang/Object;"
        映射进 super_dic
    """
    if classes.extends != "Ljava/lang/Object;":
        super_dic[class_name] = str(classes.extends)
        """
        假如他的父类在 replacemap当中，则 映射当前类 ————> 父类
        """
        if str(classes.extends) in replacemap:
            implement_dic[class_name] = str(classes.extends)

    if classes.implements:
        for imp in classes.implements:
            if str(imp) in replacemap:
                implement_dic[class_name] = str(imp)

    for method in classes.get_methods():
        if method.is_external():
            continue
        m = method.get_method()
        class_functions[class_name].append(str(m.full_name))
        c = defaultdict(int)
        flag = False
        for ins in m.get_instructions():  # count
            flag = True  # exist instructions
            c[ins.get_name()] += 1
        opcode = {}
        for p in smali_opcode:
            opcode[p] = 0
        for op in c:
            if op in smali_opcode:
                opcode[op] += c[op]
        # 将操作码向量写进CSV当中
        if flag:
            try:
                write_csv(opcode,opcodeFile, method2nodeMap[str(m.full_name)][0])
            except Exception as e:
                print(e)
opcodeFile.close()

"""----------------------------------------------PERMISSION ANALYSIS----------------------------------------------------------"""
sensitiveapimap = {}
class2runinit = defaultdict(dict)

# 从API.csv获取每个API对应的权限mapping
per_map = get_from_csv_gml('/home/chenjiaxiong/jx_code/MsDroid-JXC/feature/Andro/Data/APIs/API_all.csv')

result_f = create_csv(permission, permissionFilename)
node_attr = df_from_G(G)
getresolver = ";->getContentResolver()Landroid/content/ContentResolver;"
functions = node_attr.label
ids = node_attr.id

"""
substitude_permission
子类所涉及的敏感api的,权限列表
如 类A(自定义的) 他的父类为 B(系统自带)  ---通过这种方法可以利用原本api_permission里面的映射
则映射为 B : 权限
# start method connect run method 
class2runinit
"""
substitude_permission = substitude()

node_cp_permission = defaultdict(list)
java_class = {}  # need to generate java file

# 看看他们调用了什么权限。（就是看其后继节点）
for i in range(len(ids)):
    function = functions.get(i)
    if function.find(getresolver) >= 0:
        node_id = ids.get(i)
        # 获得当前节点的邻居
        nodes = n_neighbor(node_id, G)
        # 遍历邻居节点.
        for node in nodes:
            # 返回一个邻居的邻居的node_permission
            # deal_node  ->  处理邻居的邻居
            node_permission = deal_node(node)
            if node_permission:
                label = get_label(node, G)
                left = label.find(' ')
                right = label.find('->')
                function_class = label[left + 1: right]
                java_class.update({function_class: node_permission})
"""
getContentResolver()Landroid/content/ContentResolver （敏感函数）
java_class 以;->getContentResolver()Landroid/content/ContentResolver;为中心节点
研究其孩子节点所调用的ContentResolve方法。
"""
class_name = set ()
# 制造反编译器。

from androguard.decompiler.decompiler import DecompilerDAD
for dex in d:
    dex.set_decompiler(DecompilerDAD(dex, dx))
    for method in dex.get_methods():
        # 假如该方法的类设计到权限。
        if str(method.get_class_name()) in java_class:
            # 获得设计权限的类
            current_class = dex.get_class(method.get_class_name())
            # 获取该类的source （代码）
            content = str(current_class.get_source())
            # 返回其权限。（POP返回字典对应的内容。）
            node_permission = java_class.pop(method.get_class_name())
            # 假如有调用content的URL。 'content://ctspermissionwithsignaturepath/yes': {('R', 'Permission:com.android.cts.permissionWithSignature')}
            if content.find('content://') >= 0:
                 for per in cppermission.keys():
                        if content.find(per) >= 0:
                            pers = cppermission[per]
                            for p in pers:
                                if p[0] in node_permission:
                                    for n_id in node_permission[p[0]]:
                                        node_cp_permission[n_id].append(p[1])

# 假如节点本身涉及permission。
def count_permission(name, per_map):
        result = {}
        # 假如这个API存在于API-PER 列表中
        if name in per_map.keys():
            for p in permission:
                """
                令所有的权限都等于=0
                """
                result[p] = 0
            pers = per_map[name]
            for per in pers:
                if per not in result.keys():
                    continue
                result[per] = 1
        return result


i = 0
while i < len(ids):
    s = functions.get(i)
    s = node2function(s)
    p = count_permission(s, per_map)
    node_id = ids.get(i)
    if node_id in node_cp_permission:  # 与content provider相关的权限
        for per in permission:
            p[per] = 0
        for per in node_cp_permission[node_id]:
            p[per] = 1
    if node_id in substitude_permission:  # 子类是敏感的API
        for per in permission:
            p[per] = 0
        for per in substitude_permission[node_id]:
            p[per] = 1
    if p != {}:
        write_csv(p, result_f, node_id)
        node_permission = []
        for k in p:
            if p[k] == 1:
                node_permission.append(k)
        sensitiveapimap.update({node_id: node_permission})
    i += 1


"""----------------------------------------------TPL ANALYSIS----------------------------------------------------------"""
def get_pkg(apk_path, ratio=0.6):
    lrd = LibRadarLite(apk_path)
    pkgs = []
    try:
        res = lrd.compare()
    except Exception:
        return pkgs
    for i in res:
        p = i['Package']
        try:
            r = eval(i['Match Ratio'])
            if r >= ratio:
                pkgs.append(p)
            r = '%f' % r
        except KeyError:
            r = 'Undefined'
    return pkgs

def dfs(nodeid,deepth):
        async_baseclass = {}
        nodes = list()
        leafs = set()
        nodes.append(nodeid)
        label = get_label(nodeid,G)
        node_class = getclass(label)
        # 父节点还是用自己的？？wtf。
        # parent_class = getclass(label)
        # 如果这个类在  class2runinit 中，则更新到 async_baseclass
        if node_class in class2runinit:
            async_baseclass.update(class2runinit[node_class])
        # async_baseclass 是{45: 'java/lang/Runnable', 46: 'java/lang/Runnable'}
            # 是指这个节点是被 45 46 启动的。
        dp = 0
        # 仅仅考虑这个类不是被调用起来的。（而是调用了replacemap)
        while nodes:
            nodeid = nodes.pop()
            label = get_label(nodeid, G)
            for rk in replacemap:
                if label.find(rk) >= 0:  # 调用了start函数
                    for ck in class2runinit.keys():
                        # 如果找到了对应的类
                        if node_class.find(ck) >= 0:
                            # async类就直接添加进来。
                            async_baseclass = class2runinit[ck]
                            for asid in async_baseclass.keys():  # 查找对应的run函数
                                funcs = replacemap[rk]
                                for func in funcs:
                                    if func.find(async_baseclass[asid]) >= 0:
                                        nodes.append(asid)

            leafs.add(nodeid)
            targets = G.successors(nodeid)
            if dp < deepth:
                for t in targets:
                    if t in leafs:
                        continue
                    nodes.append(t)
            dp = dp + 1
        return leafs

TplSensitiveNodeMap = {}
tpl_list = get_pkg(apkfile)

labels = node_attr.label
ids = node_attr.id
i = 0
leafsmap = {}

def getTplSensitiveNode(nodeid):
    TplSensitiveNodes = set()
    leafs = dfs(nodeid)
    for leaf in leafs:
        if leaf in sensitiveapimap:
            TplSensitiveNodes.add(leaf)
    return TplSensitiveNodes

while i < len(ids):
    node_id = ids.get(i)
    label = labels.get(i)
    flag = False
    node_class = "L" + getclass(label)
    for tpl in tpl_list:
        if node_class.find(tpl) >= 0:
            flag = True
            break
    # 假如是第三方库。则...
    if flag:
        TplSensitiveNodeMap[node_id] = getTplSensitiveNode(node_id)
    i = i + 1